package com.guqj.zookeeper.lock;

import java.io.IOException;
import java.util.List;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

public class DistributeLock {

	private ZooKeeper zooKeeper;
	private String root;
	private LockNode idName;
	private String id; // 当前获取锁的节点ID
	private byte[] data = { 0x12, 0x34 };
	private static final int SESSION_TIMEOUT = 10000;
	private String ownerId; // 持有锁的节点ID
	private String lastChildId; // 最近的子节点

	public DistributeLock(String root, String connect) {
		try {
			this.root = root;
			this.zooKeeper = new ZooKeeper(connect, SESSION_TIMEOUT, null);
			// delAllNode(root);
			ensureExists(root);
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public void delAllNode(String path) {
		try {
			List<String> paths = zooKeeper.getChildren(root, false);
			for (String p : paths) {
				zooKeeper.delete(root + "/" + p, -1);
			}
		} catch (KeeperException e) {
			e.printStackTrace();
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
	}

	/**
	 * 判断节点是否存在，不存在则创建
	 * 
	 * @param path
	 */
	private void ensureExists(final String path) {
		try {
			Stat stat = zooKeeper.exists(path, false);
			if (stat != null) {
				return;
			}
			String result = zooKeeper.create(path, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
			System.out.println("创建节点" + path + "," + result);
		} catch (KeeperException | InterruptedException e) {
			e.printStackTrace();
		}
	}

	/**
	 * 判断当前是不是锁的owner
	 */
	public boolean isOwner() {
		return id != null && ownerId != null && id.equals(ownerId);
	}

	/**
	 * 返回当前的节点id
	 */
	public String getId() {
		return this.id;
	}

	/**
	 * 返回锁对象对应的path
	 */
	public String getRoot() {
		return root;
	}

	/**
	 * 释放锁对象
	 */
	public void unlock() {
		if (id != null) {
			try {
				zooKeeper.delete(id, -1);
				System.out.println(System.currentTimeMillis() + "-" + id + "释放锁");
			} catch (InterruptedException | KeeperException e) {
				e.printStackTrace();
			} finally {
				id = null;
			}
		}
	}

	/**
	 * 尝试获取锁对象, 不会阻塞
	 *
	 * @throws InterruptedException
	 * @throws KeeperException
	 */
	public boolean tryLock() {
		if (isOwner()) {
			return true;
		}
		acquireLock(null);
		return isOwner();
	}

	/**
	 * 尝试获取锁操作，阻塞式可被中断
	 */
	public void lock() {
		if (isOwner()) {
			return;
		}
		BooleanMutex mutex = new BooleanMutex();
		acquireLock(mutex);
		try {
			mutex.lock();
			System.out.println(System.currentTimeMillis() + "-" + id + "获取锁成功");
		} catch (Exception e) {
			e.printStackTrace();
			if (!mutex.state()) {
				lock();
			}
		}
	}

	public void lock(long timeout, TimeUnit unit) {
		if (isOwner()) {
			return;
		}
		BooleanMutex mutex = new BooleanMutex();
		acquireLock(mutex);
		try {
			mutex.lockTimeOut(timeout, unit);// 阻塞等待值为true
		} catch (Exception e) {
			e.printStackTrace();
			if (!mutex.state()) {
				lock();
			}
		}
	}

	private boolean acquireLock(final BooleanMutex mutex) {
		try {
			do {
				if (id == null) {
					long sessionId = zooKeeper.getSessionId();
					String prefix = "x-" + sessionId + "-";
					// 创建节点
					List<String> names = zooKeeper.getChildren(root, false);
					for (String name : names) {
						if (name.startsWith(prefix)) {
							id = root + "/" + name;
							break;
						}
					}
					if (id == null) {
						id = zooKeeper.create(root + "/" + prefix, data, ZooDefs.Ids.OPEN_ACL_UNSAFE,
								CreateMode.EPHEMERAL_SEQUENTIAL); // node节点为EPHEMERAL_SEQUENTIAL，自增长的临时节点
					}
					idName = new LockNode(id);
				}

				if (id != null) {
					List<String> names = zooKeeper.getChildren(root, false);
					if (names == null || names.isEmpty()) {
						id = null; // 异常情况，重新创建一个
					} else {
						// 对节点进行排序
						SortedSet<LockNode> sortedNames = new TreeSet<>();
						for (String name : names) {
							sortedNames.add(new LockNode(root + "/" + name));
						}

						if (!sortedNames.contains(idName)) {
							id = null; // 异常情况，重新创建一个
							continue;
						}

						// 将第一个节点做为ownerId
						ownerId = sortedNames.first().getId();
						if (mutex != null && isOwner()) {
							mutex.unlock();// 直接更新状态，并唤醒下一个线程
							return true;
						} else if (mutex == null) {
							return isOwner();
						}
						SortedSet<LockNode> lessThanMe = sortedNames.headSet(idName);
						if (lessThanMe != null && !lessThanMe.isEmpty()) {
							// 关注一下排队在自己之前的最近的一个节点
							LockNode lastChild = lessThanMe.last();
							lastChildId = lastChild.getId();
							// 异步watcher处理
							Stat stat = zooKeeper.exists(lastChildId, new Watcher() {
								public void process(WatchedEvent event) {
									acquireLock(mutex);
								}
							});
							if (stat != null) {
								return false;
							} else {
								acquireLock(mutex);// 如果节点不存在，需要自己重新触发一下，watcher不会被挂上去
							}
						} else {
							if (isOwner()) {
								mutex.unlock();
							} else {
								id = null;// 可能自己的节点已超时挂了，所以id和ownerId不相同
							}
						}
					}
				}
			} while (id == null);
		} catch (KeeperException e) {
			e.printStackTrace();
			if (mutex != null) {
				mutex.unlock();
			}
		} catch (InterruptedException e) {
			e.printStackTrace();
			if (mutex != null) {
				mutex.unlock();
			}
		}
		return Boolean.FALSE;
	}
}
